import random
from pygame import Rect
from pygame.constants import K_SPACE, K_LEFT, K_RIGHT, K_UP, K_DOWN
import pgzrun

piece_structures = [
    [
        [
            [" ", " ", " ", " "],
            ["i", "i", "i", "i"],
            [" ", " ", " ", " "],
            [" ", " ", " ", " "],
        ],
        [
            [" ", "i", " ", " "],
            [" ", "i", " ", " "],
            [" ", "i", " ", " "],
            [" ", "i", " ", " "],
        ],
    ],
    [
        [
            [" ", " ", " ", " "],
            [" ", "o", "o", " "],
            [" ", "o", "o", " "],
            [" ", " ", " ", " "],
        ],
    ],
    [
        [
            [" ", " ", " ", " "],
            ["j", "j", "j", " "],
            [" ", " ", "j", " "],
            [" ", " ", " ", " "],
        ],
        [
            [" ", "j", " ", " "],
            [" ", "j", " ", " "],
            ["j", "j", " ", " "],
            [" ", " ", " ", " "],
        ],
        [
            ["j", " ", " ", " "],
            ["j", "j", "j", " "],
            [" ", " ", " ", " "],
            [" ", " ", " ", " "],
        ],
        [
            [" ", "j", "j", " "],
            [" ", "j", " ", " "],
            [" ", "j", " ", " "],
            [" ", " ", " ", " "],
        ],
    ],
    [
        [
            [" ", " ", " ", " "],
            ["l", "l", "l", " "],
            ["l", " ", " ", " "],
            [" ", " ", " ", " "],
        ],
        [
            [" ", "l", " ", " "],
            [" ", "l", " ", " "],
            [" ", "l", "l", " "],
            [" ", " ", " ", " "],
        ],
        [
            [" ", " ", "l", " "],
            ["l", "l", "l", " "],
            [" ", " ", " ", " "],
            [" ", " ", " ", " "],
        ],
        [
            ["l", "l", " ", " "],
            [" ", "l", " ", " "],
            [" ", "l", " ", " "],
            [" ", " ", " ", " "],
        ],
    ],
    [
        [
            [" ", " ", " ", " "],
            ["t", "t", "t", " "],
            [" ", "t", " ", " "],
            [" ", " ", " ", " "],
        ],
        [
            [" ", "t", " ", " "],
            [" ", "t", "t", " "],
            [" ", "t", " ", " "],
            [" ", " ", " ", " "],
        ],
        [
            [" ", "t", " ", " "],
            ["t", "t", "t", " "],
            [" ", " ", " ", " "],
            [" ", " ", " ", " "],
        ],
        [
            [" ", "t", " ", " "],
            ["t", "t", " ", " "],
            [" ", "t", " ", " "],
            [" ", " ", " ", " "],
        ],
    ],
    [
        [
            [" ", " ", " ", " "],
            [" ", "s", "s", " "],
            ["s", "s", " ", " "],
            [" ", " ", " ", " "],
        ],
        [
            ["s", " ", " ", " "],
            ["s", "s", " ", " "],
            [" ", "s", " ", " "],
            [" ", " ", " ", " "],
        ],
    ],
    [
        [
            [" ", " ", " ", " "],
            ["z", "z", " ", " "],
            [" ", "z", "z", " "],
            [" ", " ", " ", " "],
        ],
        [
            [" ", "z", " ", " "],
            ["z", "z", " ", " "],
            ["z", " ", " ", " "],
            [" ", " ", " ", " "],
        ],
    ],
]



import board

CELL_SIZE = 20
X_COUNT, Y_COUNT = 10, 18

BACK_COLOR = (255, 255, 255)


class TetrisBoard(board.Board):
    """ 俄罗斯方块 """
    name = "俄罗斯方块"
    grid = []
    sequence = []
    piece_x = 3
    piece_y = 0
    piece_type = 0
    piece_rotation = 0

    def __init__(self):
        super().__init__(CELL_SIZE, X_COUNT, Y_COUNT)
        self.ext_width = self.cell_size * 4
        self.ext_height = self.cell_size * 7
        self.piece_size = 4
        self.timer_limit = 0.5
        self.timer = 0

    def reset(self):
        """ 重开游戏 """
        self.grid = []
        for y in range(self.row_count):
            self.grid.append([])
            for x in range(self.col_count):
                self.grid[y].append(" ")
        self.timer = 0
        self.new_sequence()
        self.new_piece()

    def new_sequence(self):
        self.sequence = list(range(len(piece_structures)))
        random.shuffle(self.sequence)

    def new_piece(self):
        self.piece_x = 3
        self.piece_y = 0
        self.piece_type = self.sequence.pop()
        if len(self.sequence) == 0:
            self.new_sequence()
        self.piece_rotation = 0

    def can_piece_move(self, test_x, test_y, test_rotation):
        for y in range(self.piece_size):
            for x in range(self.piece_size):
                test_block_x = test_x + x
                test_block_y = test_y + y
                if (
                    piece_structures[self.piece_type][test_rotation][y][x] != " "
                    and (
                        test_block_x < 0
                        or test_block_x >= self.col_count
                        or test_block_y >= self.row_count
                        or self.grid[test_block_y][test_block_x] != " "
                    )
                ):
                    return False
        return True


    def update(self, dt):
        self.timer += dt
        if self.timer >= self.timer_limit:
            self.timer = 0
    
            test_y = self.piece_y + 1
            if self.can_piece_move(self.piece_x, test_y, self.piece_rotation):
                self.piece_y = test_y
            else:
                # Add piece to inert
                for y in range(self.piece_size):
                    for x in range(self.piece_size):
                        block = piece_structures[self.piece_type][self.piece_rotation][y][x]
                        if block != " ":
                            self.grid[self.piece_y + y][self.piece_x + x] = block
    
                # Find complete rows
                for y in range(self.row_count):
                    complete = True
                    for x in range(self.col_count):
                        if self.grid[y][x] == " ":
                            complete = False
                            break
    
                    if complete:
                        for ry in range(y, 1, -1):
                            for rx in range(self.col_count):
                                self.grid[ry][rx] = self.grid[ry - 1][rx]

                        for rx in range(self.col_count):
                            self.grid[0][rx] = " "
    
                self.new_piece()
    
                if not self.can_piece_move(self.piece_x, self.piece_y, self.piece_rotation):
                    self.reset()
    
    
    def on_pressed(self, key):
        if key == K_UP or key == K_SPACE:
            test_rotation = self.piece_rotation + 1
            if test_rotation > len(piece_structures[self.piece_type]) - 1:
                test_rotation = 0
    
            if self.can_piece_move(self.piece_x, self.piece_y, test_rotation):
                self.piece_rotation = test_rotation
    
        elif key == K_DOWN:
            while self.can_piece_move(self.piece_x, self.piece_y + 1, self.piece_rotation):
                self.piece_y += 1
                self.timer = self.timer_limit
    
        elif key == K_LEFT:
            test_x = self.piece_x - 1
    
            if self.can_piece_move(test_x, self.piece_y, self.piece_rotation):
                self.piece_x = test_x
    
        elif key == K_RIGHT:
            test_x = self.piece_x + 1
    
            if self.can_piece_move(test_x, self.piece_y, self.piece_rotation):
                self.piece_x = test_x


    def draw_screen(self, screen):
        screen.fill(BACK_COLOR)
    
        def draw_block(block, x, y):
            colors = {
                " ": (222, 222, 222),
                "i": (120, 195, 239),
                "j": (236, 231, 108),
                "l": (124, 218, 193),
                "o": (234, 177, 121),
                "s": (211, 136, 236),
                "t": (248, 147, 196),
                "z": (169, 221, 118),
                "preview": (190, 190, 190),
            }
            color = colors[block]

            block_draw_size = self.cell_size - 1
            screen.draw.filled_rect(
                Rect(
                    x * self.cell_size, y * self.cell_size,
                    block_draw_size, block_draw_size
                ),
                color=color
            )
    
        offset_x = 2
        offset_y = 5
    
        for y in range(self.row_count):
            for x in range(self.col_count):
                draw_block(self.grid[y][x], x + offset_x, y + offset_y)
    
        for y in range(self.piece_size):
            for x in range(self.piece_size):
                block = piece_structures[self.piece_type][self.piece_rotation][y][x]
                if block != " ":
                    draw_block(
                        block,
                        x + self.piece_x + offset_x,
                        y + self.piece_y + offset_y
                    )
    
        for y in range(self.piece_size):
            for x in range(self.piece_size):
                block = piece_structures[self.sequence[-1]][0][y][x]
                if block != " ":
                    draw_block("preview", x + 5, y + 1)



game = TetrisBoard()
TITLE = game.name  # 窗口标题
WIDTH, HEIGHT = game.screen_size

def on_key_down(key):
    game.on_pressed(key)

def update(dt):
    game.update(dt)

def draw():
    screen.clear()  # 清除屏幕内容
    game.draw_screen(screen)

pgzrun.go()
